from cmath import sqrt
from email.mime import image
import math
from types import FunctionType
from typing import Any, Dict, List, Tuple
import numpy
import cv2
from PIL import ImageFont, ImageDraw, Image
from PIL.ImageFont import FreeTypeFont
import torch
from word_vec_encoders.word_vec_encoder_base import WordVecEncoderBase


class GlyphWordVecEncoder(WordVecEncoderBase):
    """
    该编码器将字符渲染为小型图像，使用卷积网络提取特征。
    """

    def __init__(self, font_file_name:str = 'SourceHanSansSC-Regular.otf', glyph_size:int = 14, embedding_dim = 128) -> None:
        super().__init__('glyph', embedding_dim = embedding_dim)
        self.font_file_name = font_file_name
        self.font_size = glyph_size
        self.font = ImageFont.truetype(font_file_name, self.font_size)
        self.image_size = glyph_size
        self.embedding_dim = embedding_dim
        self.cnn_out_channels = 32
        self.cnn = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3),
            # 此时 map_size = 14 - 2 = 12
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2, stride=2),
            # 此时 map_size = 12 / 2 = 6
            torch.nn.Conv2d(in_channels=8, out_channels=self.cnn_out_channels, kernel_size=3),
            # 此时 map_size = 6 - 2 = 4
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2, stride=2),
            # 此时 map_size = 2
        )
        self.cnn_out_map_size = 2
        self.linear_in_dim = self.cnn_out_map_size * self.cnn_out_map_size * self.cnn_out_channels
        self.linear_middle_dim = int(math.sqrt(self.linear_in_dim * self.embedding_dim))
        self.fc = torch.nn.Sequential(
            torch.nn.Linear(self.linear_in_dim, self.embedding_dim),
            torch.nn.ReLU()
        )

    def forward(self, input_tensor:torch.Tensor, batch_size:int):
        seq_batch_size, _, _, _ = input_tensor.size()
        # 通过 CNN 提取特征
        out = self.cnn(input_tensor)
        # 将特征图展平
        out = out.reshape((seq_batch_size, self.cnn_out_map_size * self.cnn_out_map_size * self.cnn_out_channels))
        # 通过全连接得到目标尺寸的嵌入
        out = self.fc(out)
        # 重新建立句子界限
        sentence_length = seq_batch_size // batch_size
        out = out.reshape((batch_size, sentence_length, self.embedding_dim))
        return out

    def collate_batch_tensor(self, batch_sentences:List[str]):
        input_tensor = self.batch_str_to_input_tensor_bert_style_cached(batch_sentences)
        batch_size, sentence_length, _, _ = input_tensor.size()        
        # 打破句子界限，组合为小序列的批次
        seq_batch_size = batch_size * sentence_length
        input_tensor = input_tensor.reshape((seq_batch_size, self.image_size, self.image_size, 1))
        # Conv2d所需的维度顺序为(seq_batch_size, n_channels, w, h)
        input_tensor = input_tensor.swapdims(2, 3)
        input_tensor = input_tensor.swapdims(1, 2)
        return input_tensor

    def sentence_to_tensor(self, sentence:str) -> torch.Tensor:
        sentence_tensor = []
        for char in sentence:
            sentence_tensor.append(self.encode_char_cached(char, self.image_size, self.font))
        return torch.stack(sentence_tensor)

    def empty_chars_tensor(self, num_chars:int) -> torch.Tensor:
        return torch.zeros((num_chars, self.image_size, self.image_size))

    _glyph_cache:Dict[str, torch.Tensor] = dict()
    def encode_char_cached(self, char:str, image_size:int, font:FreeTypeFont):
        if char not in GlyphWordVecEncoder._glyph_cache:
            image = Image.fromarray(numpy.zeros((image_size, image_size), numpy.uint8))
            draw = ImageDraw.Draw(image)
            # w, h = font.getsize(char)
            # x = (image_size - w) / 2
            # y = (image_size - h) / 2
            # 让字位于中间
            draw.text((0, -5),  char, fill ="white", font = font, align ="left")
            GlyphWordVecEncoder._glyph_cache[char] = torch.Tensor(numpy.array(image)) / 255.0
        glyph = GlyphWordVecEncoder._glyph_cache[char]
        return glyph